import torch.nn as nn


from timm.models.efficientnet import tf_efficientnetv2_l_in21k

__all__ = ['tf_efficientnetv2_l_in21k']


class EfficientNetv2(nn.Module):

    def __init__(self, num_classes, drop_rate, pretrained):
        super(EfficientNetv2, self).__init__()

        self.efficientnetv2 = tf_efficientnetv2_l_in21k(num_classes=num_classes, drop_rate=drop_rate, pretrained=pretrained)


    def forward(self, x):

        x = self.efficientnetv2(x)
        return x


def efficientnetv2(num_classes=2, drop_rate=0., pretrained=False):
    model = EfficientNetv2(num_classes, drop_rate, pretrained)
    model.default_cfg = model.efficientnetv2.default_cfg
    return model
